package cn.ichiva.dl4j.gnn.mnist;

import cn.ichiva.dl4j.gnn.common.Evolution;
import cn.ichiva.dl4j.gnn.common.UI;
import cn.ichiva.dl4j.gnn.common.Warehouse;
import cn.ichiva.dl4j.gnn.mnist.data.MnistData;
import cn.ichiva.dl4j.gnn.mnist.data.PathMnistDataSetIterator;
import lombok.extern.slf4j.Slf4j;
import org.deeplearning4j.eval.Evaluation;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;

import java.io.IOException;
import java.util.Iterator;
import java.util.Map;
import java.util.TreeMap;

@Slf4j
public class MnistEnv {

    private final TreeMap<Double, MnistAgent> list = new TreeMap<>();
    private final int max;

    /**
     * 初始化环境
     *
     * 1.向环境中初始化ai
     * 2.将初始化ai进行环境适应性测试,并排序
     * @param max
     */
    public MnistEnv(int max){
        this.max = max;

        for (int i = 0; i < max; i++) {
            MnistAgent agent = MnistAgent.newInstance();
            test(agent);
            synchronized (list) {
                list.put(agent.getScore(),agent);
            }
            log.info("初始化智能体 name = {} , score = {}",i,agent.getScore());
        }
    }

    private void test(MnistAgent agent) {
        MultiLayerNetwork network = agent.getMultiLayerNetwork();
        PathMnistDataSetIterator dataIterator = MnistData.localTestData.get();
        Evaluation eval = new Evaluation(agent.getConfig().getOutput());
        try {
            while (dataIterator.hasNext()) {
                DataSet data = dataIterator.next();
                INDArray output = network.output(data.getFeatures(), false);
                eval.eval(data.getLabels(),output);
            }
        }finally {
            dataIterator.reset();
        }
        agent.setScore(eval.accuracy());
    }

    /**
     * 进化
     *
     * 每轮随机创建ai并放入环境中进行优胜劣汰
     * @param n 进化次数
     */
    public void evolution(int n){
//        BlockThreadPool blockThreadPool=new BlockThreadPool(4);
//        for (int i = 0; i < n; i++) {
//            blockThreadPool.execute(() -> contend(newLive()));
//        }

        for (int i = 0; i < n; i++) {
            MnistAgent angent = newLive();
            contend(angent);
        }
    }

    /**
     * 随机生成新智能体
     *
     * 完全随机产生母本
     * 随机从比目标相同或更高评分中选择父本
     *
     * 基因进化在1%~10%之间进行,评分越高基于越稳定
     */
    public MnistAgent newLive(){
        double r = Math.random();
        //基因突变率
        double v = r / 11 + 0.01;
        //母本
        MnistAgent mAgent = getMother(r);
        //父本
        MnistAgent fAgent = getFather(r);

        int i = (int) (Math.random() * 3);
        INDArray newDNA = Evolution.defInstance.inheritance(mAgent.getDna(), fAgent.getDna());
        switch (i){
            case 0:
                newDNA = Evolution.defInstance.other(newDNA,v);
                break;
            case 1:
                newDNA = Evolution.defInstance.mutation(newDNA,v,0.1);
                break;
            case 2:
                newDNA = Evolution.defInstance.substitution(newDNA,v);
                break;
        }
        return MnistAgent.create(newDNA);
    }

    /**
     * 父本只选择比母本评分高的样本
     * @param r
     * @return
     */
    private MnistAgent getFather(double r) {
        r += (Math.random() * (1-r));
        return getMother(r);
    }

    private MnistAgent getMother(double r) {
        int index = (int) (r * max);
        return getMnistAgent(index);
    }

    private MnistAgent getMnistAgent(int index) {
        synchronized (list) {
            Iterator<Map.Entry<Double, MnistAgent>> it = list.entrySet().iterator();
            for (int i = 0; i < index; i++) {
                it.next();
            }

            return it.next().getValue();
        }
    }

    /**
     * 竞争
     * @param agent
     */
    public void contend(MnistAgent agent){
        test(agent);
        double score = agent.getScore();
        boolean mvp = false;
        synchronized (list) {
            if(score <= list.firstKey()){
                UI.put("无法生存",String.format("name = %s,  score = %s", agent.getName(),agent.getScore()));
                return;
            }

            if(score > list.lastKey()){
                mvp = true;
            }

            list.put(score,agent);
            if (list.size() > max) {
                MnistAgent lastAI = list.remove(list.firstKey());
                UI.put("淘 汰 ",String.format("name = %s, score = %s", lastAI.getName(),lastAI.getScore()));
            }
        }

        if(mvp){
            try {
                validation(agent);
                UI.put("MVP验证",String.format("score = %s,validScore = %s",agent.getScore(),agent.getValidScore()));
                Warehouse.write(agent);
            } catch (IOException ex) {
                log.error("保存对象失败",ex);
            }
        }
    }

    /**
     * 评估
     *
     * @param ai
     */
    public void validation(MnistAgent ai){
        MultiLayerNetwork network = ai.getMultiLayerNetwork();
        PathMnistDataSetIterator dataIterator = MnistData.localTrainData.get();
        Evaluation eval = new Evaluation(ai.getConfig().getOutput());
        try {
            while (dataIterator.hasNext()) {
                DataSet data = dataIterator.next();
                INDArray output = network.output(data.getFeatures(), false);
                eval.eval(data.getLabels(),output);
            }
        }finally {
            dataIterator.reset();
        }
        ai.setValidScore(eval.accuracy());
    }
}
